package com.phoenix.shuaidatabase.utils;


import com.phoenix.shuaidatabase.utils.ShuaiRateLimiter;

import java.util.ArrayList;
import java.util.concurrent.ConcurrentHashMap;

public class TokenBucketRateLimiter implements ShuaiRateLimiter {

    ConcurrentHashMap<String,TokenBucket> resources;

    public TokenBucketRateLimiter(ArrayList<String> resourceNames,int n) {
        this.resources = new ConcurrentHashMap<>();
        for(String resource : resourceNames) {
            resources.put(resource,new TokenBucket(n));
        }
    }

    @Override
    public boolean tryAcquire(String resource) {
        return resources.get(resource).tryAcquire();
    }

    private static class TokenBucket {
        private final int capacity; //总令牌数
        private int tokens; //目前可用令牌
        private long lastRefillTime; //上一次添加令牌的时间
        private int rate; //放入令牌的速度

        public TokenBucket(int n) {
            this.rate = n;
            this.capacity = n;
            this.tokens = n;
            this.lastRefillTime = System.currentTimeMillis();
        }

        public synchronized boolean tryAcquire() {
            //添加令牌
            refill();
            if(tokens > 0) {
                tokens --;
                return true;
            }else {
                return false;
            }
        }

        public void refill() {
            long now = System.currentTimeMillis();
            if(now - lastRefillTime > 1000) {
                int newTokens = (int)(now - lastRefillTime) / 1000 * rate;
                tokens = Math.min(tokens + newTokens,capacity);
                lastRefillTime = now;
            }
        }
    }

}
